#!/usr/bin/env python

import sys
import os

from PyQt4.QtCore import *
from PyQt4.QtGui import *

from numpy import array, linalg, dot, matrix
from math import sqrt, asin, acos, sin, cos, tan, atan, pi

colormap = {
    'a':QColor(255,0,0),
    'b':QColor(0,255,0),
    'c':QColor(0,0,255),
    'd':QColor(255,255,0),
}

RADIUS = 200

def find_z(posx, posy):
    posz = sqrt(RADIUS**2 - posx**2 - posy**2)
    return posz

def from_projective(points):
    norm = sqrt(sum([x**2 for x in points]))
    new_points = [(x/norm)*RADIUS for x in points]
    return new_points

def inside_circle(x,y, radius=RADIUS):
    return x**2 + y**2 < radius**2

def create_matrix(points):
    mat = array([points[0][0], points[1][0], points[2][0],
                    points[0][1], points[1][1], points[2][1],
                    points[0][2], points[1][2], points[2][2]])
    mat.shape = 3,3
    return mat

def create_vector(points):
    vector = array([points[0], points[1], points[2]])
    vector.shape = 3,1

    return vector

class AnchorPoint(QGraphicsEllipseItem):
    def __init__(self, x, y, code, *args):
        QGraphicsEllipseItem.__init__(self, 0, 0, 5, 5, *args)
        self.code = chr(ord('a')+code)
        self.setBrush(colormap[self.code])
        self.setPos(x,y)
        
    def mousePressEvent(self, evt):
        pass

    def mouseMoveEvent(self, evt):
        self.setPos(evt.scenePos())

class SourceImageItem(QGraphicsPixmapItem):
    def __init__(self, filename=None, *args):
        QGraphicsPixmapItem.__init__(self, *args)
        self.setImage(filename)
        
    def setImage(self, filename):
        if filename is None:
            self.buffer = QPixmap(400,400)
            self.buffer.fill(QColor(0,0,0))
        else:
            self.filename = filename
            b1 = QPixmap(filename)

            #self.buffer = b1.scaled(400,400, Qt.IgnoreAspectRatio)
            # Comment the line above and uncomment the lines below to switch
            # to a squared image inside the circle
            side = int(400/sqrt(2))
            self.imbuffer = b1.scaled(side,side, Qt.IgnoreAspectRatio)
            self.buffer = QPixmap(400,400)
            self.buffer.fill(QColor(0,0,0))

            impic = self.imbuffer.toImage()
            imbuf = self.buffer.toImage()
            for y in range(0, side):
                for x in range(0, side):
                    imbuf.setPixel(x+200-200*sqrt(2)/2, y+200-200*sqrt(2)/2, impic.pixel(x,y))
            self.buffer = QPixmap.fromImage(imbuf)
            


        self.mask = QRegion(0, 0, 400, 400, QRegion.Ellipse)
        tmp = QPixmap('mask.png')
        mask = QBitmap(tmp)
        self.buffer.setMask(mask)
        self.setPixmap(self.buffer)

    def setPixmap(self, pixmap):
        QGraphicsPixmapItem.setPixmap(self, pixmap)
        self.buffer = pixmap

class ImageScene(QGraphicsScene):
    def __init__(self, filename=None, *args):
        QGraphicsScene.__init__(self, *args)
        self.pixmapitem = SourceImageItem(filename)
        self.pixmapitem.setOffset(0,0)

        self.addItem(self.pixmapitem)
        self.points = []
        self.pixmapitem.scenePos()

    def mouseDoubleClickEvent(self, evt):
        if len(self.points) >= 4:
            return

        print "Double click at:", evt.scenePos()
        x, y = evt.scenePos().x(), evt.scenePos().y()

        if (x-200)**2+(y-200)**2 > 40000:
            return

        print "creating points"
        item = AnchorPoint(x, y, len(self.points))
        item.setZValue(len(self.points)*10+1)
        self.points.append(item)
        self.addItem(item)

    def openImage(self, filename):
        self.clear()
        self.pixmapitem.setImage(filename)
        self.update()

    def clear(self):
        self.pixmapitem.setImage(None)
        for point in self.points:
            self.removeItem(point)
            del point
        self.points = []
        self.update()
    
    def setPixmap(self, pixmap):
        self.pixmapitem.setPixmap(pixmap)



class CentralWidget(QWidget):
    def __init__(self, filename, *args):
        QWidget.__init__(self, *args)
        self.srcscene = ImageScene(filename)
        self.srcview = QGraphicsView(self.srcscene)
        self.srcview.centerOn(self.srcscene.pixmapitem)
        self.srcview.resize(400,400)

        self.tgtscene = ImageScene()
        self.tgtview = QGraphicsView(self.tgtscene)
        self.tgtview.centerOn(self.tgtscene.pixmapitem)
        self.tgtview.resize(400,400)

        layout = QHBoxLayout()
        layout.addWidget(self.srcview)
        layout.addWidget(self.tgtview)
        self.setLayout(layout)

    def toolbar_slot(self, action):
        text = action.text()

        if text == 'Draw':
            print 'starting to process'

            #Create pixel buffers
            self.tgtscene.pixmapitem.setImage(None)
            srcimg = self.srcscene.pixmapitem.buffer.toImage()
            tgtimg = self.tgtscene.pixmapitem.buffer.toImage()

            srcpoints = self.srcscene.points
            tgtpoints = self.tgtscene.points

            source = []
            target = []
            for src,tgt in zip(srcpoints, tgtpoints):
                pos = src.scenePos()
                posz = find_z(pos.x()-RADIUS, -(pos.y()-RADIUS))
                source.append((pos.x()-RADIUS, -(pos.y()-RADIUS), posz))

                pos = tgt.scenePos()
                posz = find_z(pos.x()-RADIUS, -(pos.y()-RADIUS))
                target.append((pos.x()-RADIUS, -(pos.y()-RADIUS), posz))

            print "Source points: ", source
            print "Target points: ", target

            source_matrix = create_matrix(source)
            target_matrix = create_matrix(target)

            print 'source_matrix'
            print source_matrix
            print 'target matrix'
            print target_matrix

            source_d = create_vector(source[-1])
            target_d = create_vector(target[-1])

            print "source d"
            print source_d
            print "target d"
            print target_d

            sol_src = linalg.solve(source_matrix, source_d)
            sol_tgt = linalg.solve(target_matrix, target_d)

            print 'solution source'
            print sol_src
            print 'solution target'
            print sol_tgt

            x = source_matrix*sol_src.transpose()
            y = target_matrix*sol_tgt.transpose()
            
            print 'ajuste de representantes'
            print x
            print 'ajuste de representantes tgt'
            print y

            x = linalg.inv(x)
            print 'inversa do ajuste src'
            print x
            
            transform = dot(y,x)

            print 'transform'
            print transform
            print 'transform * source d'
            a = dot(transform, source_matrix[:,0])
            b = dot(transform, source_matrix[:,1])
            c = dot(transform, source_matrix[:,2])
            print "a"
            print a#/a[2]
            print "b"
            print b#/b[2]
            print "c"
            print c#/c[2]
            print dot(transform, source_d)

            transform = linalg.inv(transform)

            # Do something
            for y in range(srcimg.height()):
                for x in range(srcimg.width()):
                    if not inside_circle(x-RADIUS,-(y-RADIUS)):
                        continue
                    orig = create_vector((x-RADIUS, -(y-RADIUS),
                                         find_z(x-RADIUS, -(y-RADIUS))))
                    new = dot(transform, orig)
                    #new = new/new[2]
            #print new
                    new = from_projective(new)
                    newx, newy = new[0], new[1]
                    newy *= -1

                    if new[2] < 0:
                        newx *= -1
                        newy *= -1

                    newx += RADIUS
                    newy += RADIUS
                    if newx >= 0 and newx < srcimg.width() and newy >= 0 and \
                                newy < srcimg.height():
                    #if inside_circle(newx+RADIUS, newy+RADIUS):
                        tgtimg.setPixel(x, y, srcimg.pixel(newx,newy))
                #print newx, " ", newy


            # Save back
            self.tgtscene.setPixmap(QPixmap.fromImage(tgtimg))
            self.tgtscene.update()

        elif text == 'Open...':
            filename = QFileDialog.getOpenFileName(None,"Open Image",
            os.getcwd(), 'Image Files (*.png *.jpg *bmp)')
            self.srcscene.openImage(filename)
            self.tgtscene.clear()

def main():
    app = QApplication(sys.argv)
    
    window = QMainWindow()

    toolbar = QToolBar()
    toolbar.addAction('Open...')
    toolbar.addAction('Draw')
    window.addToolBar(toolbar)

    try:
        filename = sys.argv[1]
    except:
        filename = None
    widget = CentralWidget(filename)
    window.setCentralWidget(widget)
    window.resize(900,500)

    QObject.connect(toolbar, SIGNAL('actionTriggered(QAction *)'),
                    widget.toolbar_slot)

    window.show()
    app.exec_()

if __name__ == '__main__':
    main()
